import heapq
from enum import Enum, auto
from dataclasses import dataclass, field
from dataclasses import dataclass
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy as np
from scipy.ndimage import map_coordinates
from environment import Environment
from math_utils import Vector2, Vector2Int, timed


@dataclass
class AlgorithmResult:
    '''A* 算法结果'''
    path: list[Vector2]
    length: float
    dose: float
    g: float
    init_time: float
    calc_time: float
    searched_ratio: float


@dataclass(order=True)
class AlgorithmNode:
    '''用于堆排序的坐标和代价类, g 仅用于堆的排序'''
    g: float
    coord: Vector2Int = field(compare=False)


class AlgorithmType(Enum):
    '''
    Attributes:
        DIJKSTRA_NO_DOSE: dijkstra 无辐射
        DIJKSTRA_DOSE: dijkstra 有辐射
        ASTAR_NO_DOSE: A* 无辐射
        ASTAR_DOSE_BASIC: A* 有辐射, 最初始的计算方法
        ASTAR_DOSE_WEIGHTED_F: A* 有辐射, f 计算时减少 g 的权重
            - 'A modified A* algorithm for path planning in the radioactive environment of nuclear facilities'
        ASTAR_DOSE_MORE_GRID: A* 有辐射, 扩大启发函数的检测范围
            - 'An improved A* algorithm for searching the minimum dose path in nuclear facilities'
        ASTAR_DOSE_RISK_MAP: A* 有辐射, 使用低分辨率预处理, 生成风险图
    '''
    DIJKSTRA_NO_DOSE = auto()
    '''dijkstra 无辐射'''
    DIJKSTRA_DOSE = auto()
    '''dijkstra 有辐射'''
    ASTAR_NO_DOSE = auto()
    '''A* 无辐射'''
    ASTAR_DOSE_BASIC = auto()
    '''A* 有辐射, 最初始的计算方法'''
    ASTAR_DOSE_WEIGHTED_F = auto()
    '''A* 有辐射, f 计算时减少 g 的权重'''
    ASTAR_DOSE_MORE_GRID = auto()
    '''A* 有辐射, 扩大启发函数的检测范围'''
    ASTAR_DOSE_RISK_MAP = auto()
    '''A* 有辐射, 使用低分辨率预处理图, 生成风险图'''


class AStar:
    '''使用 A* 算法进行寻路的类'''

    def __init__(self, environment: Environment) -> None:
        '''初始化环境'''
        self.__dirs = [
            Vector2Int(0, 1),
            Vector2Int(1, 0),
            Vector2Int(0, -1),
            Vector2Int(-1, 0),
            Vector2Int(1, 1),
            Vector2Int(1, -1),
            Vector2Int(-1, 1),
            Vector2Int(-1, -1),
        ]
        # random.shuffle(self.__dirs)
        self.__env = environment

    def run(self, start: Vector2, target: Vector2, type: AlgorithmType, **kwargs) -> AlgorithmResult:
        '''
        运行 A* 算法进行寻路
        Args:
            start: 寻路起点
            target: 寻路终点
            type: 算法类型

        Returns:
            AlgorithmResult: 包含路径、长度、剂量、时间、搜索百分比的结果
        '''
        _, init_time = timed(lambda: self.__reset_params(start, target, type, **kwargs))
        result, calc_time = timed(lambda: self.__calculate())

        result.init_time = init_time
        result.calc_time = calc_time
        return result

    def __reset_params(self, start: Vector2, target: Vector2, type: AlgorithmType, **kwargs) -> None:
        '''初始化算法参数'''
        # 根据算法类型选择对应的函数
        # self.__type = type
        self.type = str(type).split('.')[-1]
        self.__g_func, self.__h_func, self.__f_func = {
            AlgorithmType.DIJKSTRA_NO_DOSE: (self.__g_distance, lambda _: 0, self.__f_add),
            AlgorithmType.DIJKSTRA_DOSE: (self.__g_dose, lambda _: 0, self.__f_add),
            AlgorithmType.ASTAR_NO_DOSE: (self.__g_distance, self.__h_distance, self.__f_add),
            AlgorithmType.ASTAR_DOSE_BASIC: (self.__g_dose, self.__h_ave_dose, self.__f_add),
            AlgorithmType.ASTAR_DOSE_WEIGHTED_F: (self.__g_dose, self.__h_ave_dose, self.__f_weighted_add),
            AlgorithmType.ASTAR_DOSE_MORE_GRID: (self.__g_dose, self.__h_moregrid_dose, self.__f_add),
            AlgorithmType.ASTAR_DOSE_RISK_MAP: (self.__g_dose, lambda c: self.__h_risk_map(c, kwargs['ASTAR_DOSE_RISK_MAP_SIZE']), self.__f_add),
        }[type]

        # 转换为网格坐标
        self.__start = (start * self.__env.grid_density).to_vector2_int()
        self.__target = (target * self.__env.grid_density).to_vector2_int()

        # 初始化数据结构
        self.__open_set: list[AlgorithmNode] = [AlgorithmNode(0, self.__start)]
        self.__closed_set: set[Vector2Int] = set()
        self.__came_from: dict[Vector2Int, Vector2Int] = {}

        # 初始化代价数组
        grid_shape = [self.__env.map_shape.x, self.__env.map_shape.y]
        self.__g = np.full(grid_shape, float('inf'), dtype=np.float32)
        self.__distance = np.full(grid_shape, float('inf'), dtype=np.float32)
        self.__dose = np.full(grid_shape, float('inf'), dtype=np.float32)

        # 设置起点代价
        self.__g[self.__start.x, self.__start.y] = 0
        self.__distance[self.__start.x, self.__start.y] = 0
        self.__dose[self.__start.x, self.__start.y] = 0

        # 绘画使用
        self.__closed_set_list: list[Vector2Int] = []

        # 特定算法的初始化
        {
            AlgorithmType.DIJKSTRA_NO_DOSE: lambda: None,
            AlgorithmType.DIJKSTRA_DOSE: lambda: None,
            AlgorithmType.ASTAR_NO_DOSE: lambda: None,
            AlgorithmType.ASTAR_DOSE_BASIC: lambda: None,
            AlgorithmType.ASTAR_DOSE_WEIGHTED_F: lambda: None,
            AlgorithmType.ASTAR_DOSE_MORE_GRID: lambda: self.init_ASTAR_DOSE_V3(0.1),
            AlgorithmType.ASTAR_DOSE_RISK_MAP: lambda: self.init_risk_map(kwargs['ASTAR_DOSE_RISK_MAP_SIZE']),
        }[type]()

    def __calculate(self) -> AlgorithmResult:
        '''执行 A* 算法主循环'''

        while self.__open_set:
            current = heapq.heappop(self.__open_set).coord

            # 到达终点
            if current == self.__target:
                self.__closed_set_list.append(current)
                # 重建路径
                self.path: list[Vector2] = []
                self.__path: list[Vector2Int] = []
                node = self.__target
                while node != self.__start:
                    self.path.append(node / self.__env.grid_density)
                    self.__path.append(node)
                    node = self.__came_from[node]
                self.path.append(self.__start / self.__env.grid_density)
                self.__path.append(self.__start)

                return AlgorithmResult(
                    path=self.path[::-1],
                    length=self.__distance[self.__target.x, self.__target.y] / self.__env.grid_density,
                    dose=self.__dose[self.__target.x, self.__target.y] / self.__env.grid_density,
                    g=self.__g[self.__target.x, self.__target.y],
                    searched_ratio=len(self.__closed_set) / self.__env.map_shape.x / self.__env.map_shape.y * 100,
                    calc_time=0,
                    init_time=0,
                )

            # 跳过已访问的节点
            if current in self.__closed_set:
                continue
            self.__closed_set.add(current)
            self.__closed_set_list.append(current)

            # 遍历8个方向
            for dir in self.__dirs:
                next = current + dir
                if not self.__env.accessible(next):
                    continue
                # 计算新的代价
                g = self.__g[current.x, current.y] + self.__g_func(current, next)
                if g < self.__g[next.x, next.y]:
                    self.__g[next.x, next.y] = g
                    self.__distance[next.x, next.y] = self.__distance[current.x, current.y] + self.__g_distance(current, next)
                    self.__dose[next.x, next.y] = self.__dose[current.x, current.y] + self.__g_dose(current, next)

                    # 更新路径和启发式值
                    self.__came_from[next] = current
                    h = self.__h_func(next)
                    f = self.__f_func(g, h, next)
                    heapq.heappush(self.__open_set, AlgorithmNode(f, next))

        return AlgorithmResult([], 0, 0, 0, 0, 0, 0)

    def __g_distance(self, coord1: Vector2Int, coord2: Vector2Int) -> float:
        '''距离作为代价函数'''
        return Vector2Int.octile_distance(coord1, coord2)

    def __g_dose(self, coord1: Vector2Int, coord2: Vector2Int) -> float:
        '''平均剂量 * 距离作为代价函数'''
        ave_dose = (self.__env.dose[coord1.x, coord1.y] + self.__env.dose[coord2.x, coord2.y]) / 2
        return Vector2Int.chebyshev_distance(coord1, coord2) * ave_dose

    def __h_distance(self, coord: Vector2Int) -> float:
        '''距离作为启发函数'''
        return Vector2Int.chebyshev_distance(coord, self.__target)

    def __h_ave_dose(self, coord: Vector2Int) -> float:
        '''平均剂量 * 距离作为启发函数'''
        ave_dose = (self.__env.dose[coord.x, coord.y] + self.__env.dose[self.__target.x, self.__target.y]) / 2
        return Vector2Int.chebyshev_distance(coord, self.__target) * ave_dose

    def __h_moregrid_dose(self, coord: Vector2Int) -> float:
        '''扩大启发函数的检测范围'''
        # 计算距离终点距离
        L = Vector2Int.chebyshev_distance(coord, self.__target)
        # 小于 L_s 时, 算法转为 Dijkstra
        if L < self.Ls:
            return 0
        # 将 L 分为 n 段
        n = 4
        # 衡量当前估算路径与标准路径 (直线路径) 之间的偏离程度
        P = 1
        # 整体剂量平均值
        R = 0
        # 获取格子的数据
        grid_size = (self.__target - coord) / n
        bias = Vector2(0, grid_size.y) if abs(grid_size.x) > abs(grid_size.y) else Vector2(grid_size.x, 0)
        for i in range(n):
            # 获取每组 3 个格子的剂量
            grid_coord = [
                (coord + grid_size * (i + 0.5) - bias).to_vector2_int(),
                (coord + grid_size * (i + 0.5)).to_vector2_int(),
                (coord + grid_size * (i + 0.5) + bias).to_vector2_int(),
            ]
            grid_dose = np.array([
                self.__env.dose[grid_coord[0].x, grid_coord[0].y] if self.__env.in_map(grid_coord[0]) else float('inf'),
                self.__env.dose[grid_coord[1].x, grid_coord[1].y] if self.__env.in_map(grid_coord[1]) else float('inf'),
                self.__env.dose[grid_coord[2].x, grid_coord[2].y] if self.__env.in_map(grid_coord[2]) else float('inf'),
            ])
            # 如果最大剂量在中间, 则表明路径偏离
            if grid_dose.argmax() == 1:
                P += 1 / n
            # 更新剂量值
            grid_dose.sort()
            R += grid_dose[0] + grid_dose[1]
        # 一共计算 2n 个格子
        R /= 2 * n
        return R * L * P

    def __h_risk_map(self, coord: Vector2Int, size: int) -> float:
        '''将风险图的线性插值当做启发函数'''
        step = self.__env.map_shape / size
        x = coord.x / step.x
        y = coord.y / step.y
        value = map_coordinates(self.risk_map, [[x], [y]], order=1, mode='nearest')
        return float(value[0])

    def __f_add(self, g: float, h: float, _: Vector2Int) -> float:
        '''将 g 与 h 相加'''
        return g + h

    def __f_weighted_add(self, g: float, h: float, current: Vector2Int) -> float:
        '''加权相加, 减少 g 的权重'''
        # 将 g 的大小加权到和 h 相近
        r = Vector2Int.euclidean_distance(current, self.__target)
        R = Vector2Int.euclidean_distance(self.__start, self.__target)
        return (1 + r / R) * g + h

    def init_ASTAR_DOSE_V3(self, ratio: float):
        '''
        初始化 ASTAR_DOSE_V3 需要的距离参数
        Args:
            size: L_s 取地图的比例
        '''
        self.Ls: float = (max(self.__env.map_shape.x, self.__env.map_shape.y) * ratio)

    def init_risk_map(self, size: int) -> None:
        '''
        生成风险图
        Args:
            size: 风险图尺寸
        '''
        # 生成仅有剂量的缩略图
        h, w = self.__env.map.shape[:2]
        row_idx = np.linspace(0, h - 1, size, dtype=int)
        col_idx = np.linspace(0, w - 1, size, dtype=int)
        dose_low_resolution = self.__env.dose[np.ix_(row_idx, col_idx)]
        target = (Vector2(
            self.__target.x / self.__env.map_shape.x,
            self.__target.y / self.__env.map_shape.y,
        ) * size).to_vector2_int()
        target = Vector2Int(
            min(size - 1, max(0, target.x)),
            min(size - 1, max(0, target.y)),
        )

        # 使用 Dijkstra 计算风险图
        open_set = [AlgorithmNode(0, target)]
        closed_set = set()
        self.risk_map = np.full([size, size], float('inf'), dtype=np.float32)
        self.risk_map[target.x, target.y] = 0
        while open_set:
            current = heapq.heappop(open_set).coord
            if current in closed_set:
                continue
            closed_set.add(current)
            for dir in self.__dirs:
                next = current + dir
                if next.x < 0 or next.x >= size or next.y < 0 or next.y >= size:
                    continue
                ave_dose = (dose_low_resolution[current.x, current.y] + dose_low_resolution[next.x, next.y]) / 2
                cost = ave_dose * Vector2Int.chebyshev_distance(current, next)
                g = self.risk_map[current.x, current.y] + cost
                if g < self.risk_map[next.x, next.y]:
                    self.risk_map[next.x, next.y] = g
                    heapq.heappush(open_set, AlgorithmNode(g, next))

        self.risk_map *= max(self.__env.map_shape.x, self.__env.map_shape.y) / size

        # 绘制风险图
        # fig = plt.figure(figsize=(6, 6))
        # if fig.canvas.manager:
        #     fig.canvas.manager.set_window_title('risk_map')
        # axes = fig.add_subplot(1, 1, 1)
        # axes.set_xlim(-0.5, size - 0.5)
        # axes.set_ylim(-0.5, size - 0.5)
        # axes.axis('off')
        # axes.imshow(self.risk_map, cmap='plasma')
        # plt.show()

    def draw(self, show=False, save_to='') -> None:
        '''绘制地图、路径和搜索区域'''
        # 生成图像
        image = self.__env.get_image()

        # 标记搜索过的区域
        for coord in self.__closed_set:
            image[coord.x, coord.y] = 0.8 * image[coord.x, coord.y] + (0.2, 0.2, 0.2)

        # 标记路径
        for coord in self.__path:
            image[coord.x, coord.y] = (1, 0, 0)

        # 创建画布
        fig = plt.figure(figsize=(6, 6))
        if fig.canvas.manager:
            fig.canvas.manager.set_window_title(self.type)
        axes = fig.add_subplot(1, 1, 1)
        axes.set_xlim(-0.5, self.__env.map_shape.x - 0.5)
        axes.set_ylim(-0.5, self.__env.map_shape.y - 0.5)
        axes.axis('off')

        axes.imshow(image)
        if save_to != '':
            plt.savefig(save_to)
        if show:
            plt.show()
        else:
            plt.close(fig)

    def animate(self, sep=1, show=False, save_to='') -> None:
        '''可视化寻路过程'''

        fig, ax = plt.subplots(figsize=(6, 6))
        if fig.canvas.manager:
            fig.canvas.manager.set_window_title(self.type)
        ax.set_xlim(-0.5, self.__env.map_shape.x - 0.5)
        ax.set_ylim(-0.5, self.__env.map_shape.y - 0.5)
        ax.axis('off')

        closed_set_length = len(self.__closed_set_list)
        im = image = None

        # 初始显示地图
        def init():
            nonlocal image, im
            image = self.__env.get_image()
            im = ax.imshow(image)
            for coord in self.__path:
                image[coord.x, coord.y] = (1, 0, 0)
            return []

        def update(frame):
            nonlocal image, im
            assert im is not None and image is not None
            coords = self.__closed_set_list[frame * sep:(frame + 1) * sep]

            # 标记搜索过的区域
            for coord in coords:
                image[coord.x, coord.y] = 0.8 * image[coord.x, coord.y] + (0.2, 0.2, 0.2)

            im.set_array(image)
            ax.set_title(f"Step: {min((frame + 1) * sep, closed_set_length)} / {closed_set_length}")
            return [im]

        self.anim = FuncAnimation(
            fig,
            init_func=init,
            func=update,
            frames=len(self.__closed_set_list) // sep + 1,
            interval=0,
            repeat=False,
        )
        if save_to != '':
            self.anim.save(save_to, writer='pillow', fps=30)
        if show:
            plt.show()
